# 导入必要的库
import json  # 用于JSON数据处理
import httpx  # 异步HTTP客户端
import sys  # 系统相关功能
import os  # 操作系统相关功能
from typing import Optional, Dict, Any, AsyncGenerator  # 类型注解
# 从项目配置导入设置
from config import settings
# 导入日志记录器
from utils.logger import get_logger

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 创建日志记录器实例
logger = get_logger(__name__)


class DoubaoModel:
    """豆包大模型调用接口"""

    def __init__(self):
        # 从设置中获取API密钥和模型名称
        self.api_key = settings.DOUBAO_API_KEY
        self.model_name = settings.DOUBAO_MODEL_NAME

        # 创建异步HTTP客户端
        self.client = httpx.AsyncClient()

        # 设置请求头
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",  # 认证头
            "Content-Type": "application/json"  # 内容类型
        }
        logger.info(f"Initialized DoubaoModel with model: {self.model_name}")

    async def generate_text(self, prompt: str, **kwargs) -> str:
        """
        生成文本（同步方式）

        Args:
            prompt: 用户输入的提示词
            **kwargs: 其他可选参数

        Returns:
            模型生成的文本
        """
        url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"

        # 构造消息格式
        messages = [{"role": "user", "content": prompt}]

        # 构造请求负载
        payload = {
            "model": self.model_name,
            "messages": messages,
            "stream": False,  # 非流式
            **kwargs  # 合并其他参数
        }

        try:
            # 记录日志（只记录前50个字符）
            logger.info(f"Calling Doubao model with prompt: {prompt[:50]}...")

            # 发送POST请求
            response = await self.client.post(
                url,
                headers=self.headers,
                json=payload,
                timeout=60.0  # 60秒超时
            )
            # 检查HTTP状态码
            response.raise_for_status()

            # 解析JSON响应
            result = response.json()
            # 提取生成的文本
            generated_text = result["choices"][0]["message"]["content"]
            logger.info("Successfully generated text with Doubao model")

            return generated_text
        except Exception as e:
            logger.error(f"Error calling Doubao model: {str(e)}")
            raise  # 重新抛出异常

    async def generate_text_stream(self, prompt: str, **kwargs) -> AsyncGenerator[str, None]:
        """
        流式生成文本

        Args:
            prompt: 用户输入的提示词
            **kwargs: 其他可选参数

        Yields:
            生成的文本片段（实时）
        """
        url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"

        messages = [{"role": "user", "content": prompt}]

        payload = {
            "model": self.model_name,
            "messages": messages,
            "stream": True,  # 流式模式
            **kwargs
        }

        try:
            logger.info(f"Calling Doubao model stream with prompt: {prompt[:50]}...")

            # 使用流式请求
            async with self.client.stream(
                    "POST",
                    url,
                    headers=self.headers,
                    json=payload,
                    timeout=30.0  # 30秒超时
            ) as response:
                response.raise_for_status()

                # 逐块读取响应
                async for chunk in response.aiter_text():
                    if chunk.startswith("data:"):
                        data = chunk[5:].strip()  # 去除"data:"前缀
                        if data != "[DONE]":  # 检查是否结束
                            try:
                                result = json.loads(data)
                                # 检查是否有有效内容
                                if "choices" in result and len(result["choices"]) > 0:
                                    content = result["choices"][0]["delta"]["content"]
                                    if content:  # 如果有内容则返回
                                        yield content
                            except Exception as e:
                                logger.warning(f"Error parsing stream data: {e}")
                                continue  # 继续处理下一个块
            logger.info("Successfully streamed text with Doubao model")
        except Exception as e:
            logger.error(f"Error calling Doubao model stream: {str(e)}")
            raise

    async def close(self):
        """关闭HTTP客户端连接"""
        await self.client.aclose()


# 全局模型实例（单例模式）
doubao_model: Optional[DoubaoModel] = None


def get_doubao_model() -> DoubaoModel:
    """获取豆包模型实例（单例模式）"""
    global doubao_model
    if doubao_model is None:
        doubao_model = DoubaoModel()  # 懒加载
    return doubao_model


async def call_doubao(prompt: str, **kwargs) -> str:
    """
    调用豆包模型生成文本（便捷函数）

    Args:
        prompt: 用户输入的提示词
        **kwargs: 其他可选参数

    Returns:
        模型生成的文本
    """
    model = get_doubao_model()  # 获取模型实例
    return await model.generate_text(prompt, **kwargs)


async def call_doubao_stream(prompt: str, **kwargs) -> AsyncGenerator[str, None]:
    """
    流式调用豆包模型生成文本（便捷函数）

    Args:
        prompt: 用户输入的提示词
        **kwargs: 其他可选参数

    Yields:
        模型实时生成的文本片段
    """
    model = get_doubao_model()
    # 使用异步生成器
    async for chunk in model.generate_text_stream(prompt, **kwargs):
        yield chunk
